# 聚合因子生成器
# file: data/feature_aggregator.py

import pandas as pd
from data.stock_data_loader import StockDataLoader


class FeatureAggregator:
    def __init__(self, data_loader=None):
        self.loader = data_loader or StockDataLoader()

    def generate_daily_aggregate_features(self, start_date=None, end_date=None) -> pd.DataFrame:
        """
        对 stock_hit_board 进行聚合，按股票代码和日期提取每日特征
        """
        df = self.loader.get_hit_board_data(start_date=start_date, end_date=end_date)

        if df.empty:
            print("[Aggregator] 输入 hits 数据为空")
            return pd.DataFrame()

        # 构建聚合函数
        agg_dict = {
            "money": ["count", "sum", "max", "mean"],
            "deal_money": ["max"],
            "calculate_max_price": ["mean"],
            "limit_up_count": ["sum"],
            "ignite": lambda x: (x == 1).sum()
        }

        grouped = df.groupby(["stock_code", "add_date"]).agg(agg_dict)
        grouped.columns = [
            "hit_count", "money_sum", "money_max", "money_mean",
            "deal_money_max", "calculate_max_price_mean",
            "limit_up_count_sum", "ignite_count"
        ]
        grouped = grouped.reset_index()

        return grouped
